package com.zrj.gateway.filter;

import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.function.Function;
import java.util.function.Supplier;

/**
 * 跟CachedBodyOutputMessage几乎一样，只是将boolean isCached()改为public boolean isCached()
 *
 * @author zrj
 * @since 2022/10/7
 **/
public class ModifyCachedBodyOutputMessage implements ReactiveHttpOutputMessage {
    private final DataBufferFactory bufferFactory;
    private final HttpHeaders httpHeaders;
    private boolean cached = false;
    private Flux<DataBuffer> body = Flux.error( new IllegalStateException( "The body is not set. Did handling complete with success?" ) );

    public ModifyCachedBodyOutputMessage(ServerWebExchange exchange, HttpHeaders httpHeaders) {
        this.bufferFactory = exchange.getResponse().bufferFactory();
        this.httpHeaders = httpHeaders;
    }

    @Override
    public void beforeCommit(Supplier<? extends Mono<Void>> action) {
    }

    @Override
    public boolean isCommitted() {
        return false;
    }

    public boolean isCached() {
        return this.cached;
    }

    @Override
    public HttpHeaders getHeaders() {
        return this.httpHeaders;
    }

    @Override
    public DataBufferFactory bufferFactory() {
        return this.bufferFactory;
    }

    public Flux<DataBuffer> getBody() {
        return this.body;
    }

    /**
     * @deprecated
     */
    @Deprecated
    public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        this.body = Flux.from( body );
        this.cached = true;
        return Mono.empty();
    }

    @Override
    public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return this.writeWith( Flux.from( body ).flatMap( (p) -> p ) );
    }

    @Override
    public Mono<Void> setComplete() {
        return this.writeWith( Flux.empty() );
    }
}
